from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains.retrieval import create_retrieval_chain
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import OllamaLLM
from langchain_text_splitters import RecursiveCharacterTextSplitter

from model.MyOllamaEmbeddings import MyOllamaEmbeddings

# 加载文档
loader = TextLoader("qa.txt", encoding="UTF-8")
docs = loader.load()

# 数据切分
text_splitter = RecursiveCharacterTextSplitter() # 文档太大需要分块，分割算法
documents = text_splitter.split_documents(docs)

# 创建embeddings
embeddings = MyOllamaEmbeddings(model="nomic-embed-text")

# 通过向量数据库存储，类似于jdbc的连接实例
vectorDB = FAISS.from_documents(docs, embeddings)
# 查询检索
# 创建prompt
prompt = ChatPromptTemplate.from_template(
    """ 仅根据提供的上下文回答以下问题：
    <context>
    {context}
    </context>
    Question:{input}""")
## 需要利用语言大模型里的总结归纳能力。
llm = OllamaLLM(model="llama3.2:3b")

# 创建document 的chain查询
## 第1个chian：create_stuff_documents_chain:  文档处理的chain
documents_chain = create_stuff_documents_chain(llm, prompt)

# 创建搜索chain，返回值为vectorStoreRetriever'
"""
多种向量数据库，通过as_retriever方法，统一封装了成了一个retriever。
"""
retriever = vectorDB.as_retriever() # 拿到向量数据库检索的对象retriever，通过这个对象就可以检索数据了。
#第2个chain：向量检索的
retrieval_chain = create_retrieval_chain(retriever, documents_chain)

# 执行请求
# response = retrieval_chain.invoke({"input": "你能帮我做什么事情"})
response = retrieval_chain.invoke({"input": "天气"})
print(response["answer"])

## 我可以帮助你查询天气！要知道当前天气情况，可以输入“天气”并询问我。或者，想要知道未来几天的天气预报吗？
## 目前的天气是晴天， temperature 为 25 度。